/*
 * To change this template, choose Tools | Templates
 * and open the template in the editor.
 */

package FMMLogic;

import be.ac.ulg.montefiore.run.jahmm.*;
import Domain.Output;
import Domain.State;
import java.util.Random;
import java.util.Arrays;
import java.util.List;
import java.util.ArrayList;
import DataAccess.HmmWriter;
import DataAccess.ObservationReader;
/**
 *
 * @author Allen
 */
public class LFMMTrainer {
    /* Input:
     *   datafile
     * Output:
     *   Hmm models for each data
     */
    private List<Hmm<ObservationDiscrete<Output>>> hmms;
    private List<Output[]> outputs;

    //hard coded state transition matrix
    static double [][] T ={
                           {0.956,0.044,0.0},
                            {0.00,0.672,0.328},
                            {0.0,0.0,1.0},
                           };

    static double [][] outProb = {
                               {0.787, 0.113, 0.1},
                               //{1.0,0.0,0.0},
                               // {0.9,0.1,0.0},
                                //{0.2,0.6,0.2},
                              {0.358,  0.226, 0.416},
                               // {0.3,0.3,0.4},
                                //{0.3,0.3,0.4},
                               {0.104, 0.338, 0.578},
                               // {0.1,0.2,0.7},
                                 };

//initalize a 2D arra - each row has equal probability
       static void uniformInit (double [][] A)
       {
           for ( int i=0 ; i < A.length ; i++ ) {
               double unitDist = 1.0/ (double) A[i].length;
               Arrays.fill(A[i], unitDist);
           }
       }
       /** used for random number generation in the perturb routines.
        * seeded to 0 to ensure determinism across runs
        *
        * @see also perturb
        */
       static Random aRandom = new Random( 997L);
       /**
        * Create a new version of A, with two
        * entires in a single row perturbed. Guaranteed
        * to preserve the invariant that the row entires sum
        * to 1, and that each entry is within [0,1]
        */
        static double [][] perturb(double [][] A, double perturbFactor)
    {
            //warning- very expensive copy, should to store
            //state to just change A, and undo if does not lead
            //to better solution
            double [][] result = new double [A.length][A[0].length];
            for (int i = 0; i < A.length;i++) {
                for (int j=0; j < A[0].length; j++) {
                    result [i][j] = A[i][j];
                }
            }
            int r = aRandom.nextInt(A.length);
            int c0 = aRandom.nextInt(A[r].length);
            int c1 = aRandom.nextInt(A[r].length);

            double l = Math.min(A[r][c0], A[r][c1]);
            double max = Math.max(A[r][c0], A[r][c1]);
            double u = 1.0 - max;
            double delta = Math.min(l, u);
            double randUnit = (aRandom.nextInt() % 2 ==1) ? 1.0 : -1.0;
            delta *= randUnit;
            /*System.out.println("Random params: "+ r + " " + c0 + " " + c1
                    + " " + randUnit );*/
            result [r][c0] += delta * perturbFactor;
            result [r][c1] -= delta * perturbFactor;
            return result;
    }
    static double computeFom(Output [] trace, double [][] T,
                                double [][] O)
    {
        int numStates = T.length;
        double [] initDist = new double [numStates];
        //all states are equally likely at first
        //double uniformDist =1.0/ (double) numState;
        //Arrays.fill(p[0], uniformDist);
        //always start in the C state;
        initDist [State.C.toInt()] = 1.0;

        double [] terms = new double [T.length];
        double [][] outprob = new double [trace.length][T.length];
        double [][] logoutprob = new double [trace.length][T.length];
        for ( int i = 0 ; i < trace.length ; i++) {
            if ( i ==0) {
                for (int j = 0; j < numStates; j++ ) {
                   outprob[i][j] = initDist [j]*O[j][trace[i].toInt()];
                   //TODO: very poor choice of names- outprob & outProb
                   //have already aliased with disastrous results once
                }
                for (int j = 0; j < numStates; j++) {
                    logoutprob [i][j] = Math.log(outprob[i][j]);
                }
              } else {
                for ( int j = 0 ; j < numStates; j++) {
                    for (int k = 0; k < numStates; k++) {
                        outprob[i][j] += outprob[i-1][k] * T[k][j]
                                * O[j][trace[i].toInt()];
                        //TODO: not sure if setting terms correctly
                        terms[k] = logoutprob [i-1][k]
                                    + Math.log(T[k][j])
                                    + Math.log(O[j][trace[i].toInt()]);
                        if ( Double.isNaN(terms [k]) ) {
                        }
                    }
                    logoutprob [i][j] = logsum (terms);
                }
              }
            }
            double totalprob = 0.0 ;
            for ( int j = 0; j < numStates; j++) {
                totalprob += outprob [trace.length-1][j];
            }
            /*System.out.println("returned from : " + totalprob);*/
            double totalProbFromLogs = 0.0;
            for ( int j = 0; j < numStates; j++) {
                totalProbFromLogs += Math.exp( logoutprob[trace.length-1][j] );
            }
            /*System.out.println("fom as computed by totalProbFromLogs : " +
                    totalProbFromLogs );*/
            return totalprob;
        }
     static double logsum( double [ ] terms) {

        double result = 0;
        int i;
        for ( i = 0 ; i < terms.length; i++ ){
            if (terms [i] != Double.NEGATIVE_INFINITY) {
                break ;
            }
        }
        if( i == terms.length) {
            return Double.NEGATIVE_INFINITY;
        }

        double mu = 0.0;
        for (i = 0; i < terms.length; i ++){
            mu = (terms[i] != Double.NEGATIVE_INFINITY)
                    && (terms[i] < mu) ? terms[i] : mu;
        }
        double tmp = 0.0;

        for(i = 0; i < terms.length; i ++) {
            tmp += Math.exp( terms[i] - mu);
        }

        result = mu + Math.log(tmp);
        return result;
    }

    static final double initPerturbFactor = 0.3;
    static final int numPerturbsPerValue = 10;
    static final double updateFactor = 2.0;
    static final double minPerturbFactor = 0.01;

    static Hmm<ObservationDiscrete<Output>> convertHmm(double [][] T,
                                                            double [][]O)
    {
        //creating new Hmm with 3 states with observed Ouput class
        Hmm<ObservationDiscrete<Output>> hmm =
			new Hmm<ObservationDiscrete<Output>>(3,
					new OpdfDiscreteFactory<Output>(Output.class));
        //setting inital state probability
	hmm.setPi(0, 1.0);
	hmm.setPi(1, 0.0);
        hmm.setPi(2, 0.0);

        //setting observation probability
        for ( int i=0; i<3; i++)
            hmm.setOpdf(i, new OpdfDiscrete<Output>(Output.class,O[i]));

        //setting transition probability
        for (int i=0; i < 3; i++)
            for (int j=0; j < 3; j++)
                hmm.setAij(i, j, T[i][j]);

        return hmm;
    }

    public List<Hmm<ObservationDiscrete<Output>>> getHmms(){
        return this.hmms;
    }
    public List<Output[]> getOutputs(){
        return this.outputs;
    }
    public Hmm<ObservationDiscrete<Output>>  learnLFMM(Output[] trace) {
        double [][] Tsolve = new double[State.cardinality()][State.cardinality()];
        double [][] Osolve = new double[State.cardinality()][Output.cardinality()];
        uniformInit(Tsolve);
        uniformInit(Osolve);

        Tsolve = T;
        Osolve = outProb;

        double perturbFactor = initPerturbFactor;
        int numPerturb = numPerturbsPerValue;
        int unsuccPerturb = 0;
        while (true) {
            double fom = computeFom(trace, Tsolve, Osolve);
            double [][] Tnew = perturb( Tsolve, perturbFactor);
            double [][] Onew = perturb( Osolve, perturbFactor);

            double newfom = computeFom(trace, Tnew, Onew);
            if(newfom > fom){
                Tsolve = Tnew;
                Osolve = Onew;
            } else {
                unsuccPerturb ++;
            }

            if( unsuccPerturb > numPerturb ) {
                perturbFactor /= updateFactor;
                unsuccPerturb = 0;
            }
            if(perturbFactor < minPerturbFactor) {
                break;
            }
        }
        Hmm<ObservationDiscrete<Output>> hmm= convertHmm(Tsolve, Osolve);
        return hmm;
    }
    public void train(String resultfile){
        for (int i = 0; i < outputs.size(); i ++){
            hmms.add(this.learnLFMM(outputs.get(i)));
        }
        HmmWriter.writeHmms(resultfile, hmms);
    }
    
    public LFMMTrainer(String datafile){
        this.outputs = ObservationReader.readOutputs(datafile);
        this.hmms = new ArrayList<Hmm<ObservationDiscrete<Output>>>();

    }
}
